#include "beamformer.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "os_support.h"
#include "fftwrap.h"

typedef struct Beamformer_
{
    short *buf;
    float **x;
    float **X;
    float *win;
    float *Y;
    float *y;
    float *outbuf;
    int sr;
    int M;
    int framelen;
    void *dft;
}Beamformer;

// generate windows
static void SqrtHanning(float *win, int len)
{
    int n; float PI = 3.1415926f;
	for (n=0; n<len/2+1; n++)
	{
		win[n] = sqrtf(0.5*(1-cos(2*PI*n/len)));
	}
	for (n=1; n<len/2; n++)
	{
		win[len/2+n] = win[len/2-n];
	}
}

void *create_beamformer(int framelen, int sr, int mic)
{
    Beamformer *self = (Beamformer *)speex_alloc(sizeof(Beamformer));

    self->framelen = framelen;
    self->M = mic;
    self->sr = sr;
    self->buf = (short *)speex_alloc(sizeof(short)*framelen*2);
    self->x = (float **)speex_alloc(sizeof(float *)*mic);
    for (int m=0; m<mic; m++) {
        self->x[m] = (float *)speex_alloc(sizeof(float)*framelen*2);
    }
    self->X = (float **)speex_alloc(sizeof(float *)*mic);
    for (int m=0; m<mic; m++) {
        self->X[m] = (float *)speex_alloc(sizeof(float)*framelen*2);
    }
    self->win = (float *)speex_alloc(sizeof(float)*framelen*2);
    self->Y = (float *)speex_alloc(sizeof(float)*framelen*2);
    self->y = (float *)speex_alloc(sizeof(float)*framelen*2);
    SqrtHanning(self->win, framelen*2);
    self->dft = spx_fft_init(framelen*2);
    self->outbuf = (float *)speex_alloc(sizeof(float)*framelen);
    return (void *)self;
}

void process_beamformer(void *obj, short *ch_x, short *bf_y)
{
    Beamformer *self = (Beamformer *)obj;
    /**/
    for (int m=0; m<self->M; m++) {
        memcpy(self->x[m], self->x[m] + self->framelen, sizeof(float)*self->framelen);
        for (int i=0; i<self->framelen; i++) {
            self->x[m][self->framelen+i] = (float)ch_x[self->M*i+m];
        }
        for (int n=0; n<self->framelen*2; n++) {
            self->y[n] = self->x[m][n]*self->win[n];
        }
        spx_fft_float(self->dft, self->y, self->X[m]);               
    }
    /**/
    memset(self->Y, 0, sizeof(float)*self->framelen*2);
    self->Y[0] = 0.0f; // (self->X[0][0] + self->X[1][0])/2.0f;
    self->Y[2*self->framelen-1] = 0.0f;
    for (int nf=1; nf<self->framelen; nf++) {
        self->Y[2*nf-1] = (self->X[0][2*nf-1] + self->X[1][2*nf-1])/2.0f;
        self->Y[2*nf] = (self->X[0][2*nf] + self->X[1][2*nf])/2.0f;
    }
    spx_ifft_float(self->dft, self->Y, self->y);
    /**/
    float tmp;
    for (int n=0; n<self->framelen; n++) {
        tmp = self->y[n]*self->win[n] + self->outbuf[n];
        bf_y[n] = (short)tmp;
    }
    for (int n=self->framelen; n<self->framelen*2; n++) {
        self->outbuf[n-self->framelen] = self->y[n]*self->win[n];
    }
    
}

void destroy_beamformer(void *obj)
{
    Beamformer *self = (Beamformer *)obj;
    if (self) {
        for (int m=0; m<self->M; m++) {
            speex_free(self->x[m]);
            speex_free(self->X[m]);        
        }
        speex_free(self->x);
        speex_free(self->X);
        speex_free(self->buf);
        speex_free(self->win);
        speex_free(self->Y);
        speex_free(self->y);
        speex_free(self->outbuf);
        spx_fft_destroy(self->dft);
        speex_free(self);
    }
}